import os

from PyPDF2 import PdfFileReader, PdfFileWriter


def pdf_partition(file_path: str, not_include_pages: [], output_dir: str = None):
    """
    分割pdf

    :param file_path: .pdf文件路径
    :param not_include_pages: 不包含的页码的列表
    如果是二维
    如：[[1,4], [5,8]]表示页码除了从1到4,从5到8以外，剩下的合并输出
    如果是一维
    如：[1,4,5]表示除了1，4，5页，将源pdf的其他的剩下的页合并输出
    :param output_dir: 输出文件夹
    :return:
    """
    if output_dir is None:
        output_dir = os.path.abspath(os.path.join(file_path, *(['..'] * 1)))
    pdf_name = (file_path[len(output_dir) + 1:]).split('.')[0]
    pdf_reader = PdfFileReader(file_path)
    if type(not_include_pages[0]) == list:
        not_include_pages_set = set()
        for i in not_include_pages:
            not_include_pages_set.update(set(range(i[0], i[1] + 1)))
        not_include_pages = not_include_pages_set
    else:
        not_include_pages = set(not_include_pages)
    output_file = '{}/{}_output.pdf'.format(output_dir, pdf_name)
    if os.path.exists(output_file):
        raise Exception('输出文件[{}]已经存在！'.format(output_file))
    pdf_writer = PdfFileWriter()
    for page in range(pdf_reader.getNumPages()):
        if page + 1 not in not_include_pages:
            pdf_writer.addPage(pdf_reader.getPage(page))
    print('完成！输出文件[{}]！'.format(output_file))
    with open(output_file, "wb") as f:
        pdf_writer.write(f)


def merge_pdfs(pdf_path1: str, pdf_path2: str, output_file: str = None):
    """
    合并两个PDF文件

    :param pdf_path1: 第一个.pdf文件路径
    :param pdf_path2: 第二个.pdf文件路径
    :param output_dir: 输出文件夹，如果未指定，则默认保存在第一个PDF文件的同级目录下
    :return:
    """
    if output_file is None:
        output_dir = os.path.abspath(os.path.join(pdf_path1, *(['..'] * 1)))
        output_file = os.path.join(output_dir, 'merged_output.pdf')

    if os.path.exists(output_file):
        raise Exception(f'输出文件[{output_file}]已经存在！')

    pdf_writer = PdfFileWriter()

    # 读取并添加第一个PDF的所有页面
    pdf_reader1 = PdfFileReader(pdf_path1)
    for page in range(pdf_reader1.getNumPages()):
        pdf_writer.addPage(pdf_reader1.getPage(page))

    # 读取并添加第二个PDF的所有页面
    pdf_reader2 = PdfFileReader(pdf_path2)
    for page in range(pdf_reader2.getNumPages()):
        pdf_writer.addPage(pdf_reader2.getPage(page))

    # 将合并后的PDF写入文件
    with open(output_file, "wb") as f:
        pdf_writer.write(f)

    print(f'PDF合并完成，文件保存在：{output_file}')
